#include "server/websocket_server.h"
#include "server/websocket_session.h"
#include "client/websocket_client.h"
#include <iostream>
#include <thread>
#include <chrono>
#include <cassert>
#include <vector>

void test_basic_communication() {
    std::cout << "测试1: 基本通信测试..." << std::endl;
    
    // 启动服务器
    websocket::WebSocketServer server(8081);
    server.start();
    
    // 等待服务器启动
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    
    // 创建客户端
    websocket::WebSocketClient client("127.0.0.1", 8081);
    
    bool message_received = false;
    client.setMessageHandler([&message_received](const std::string& msg) {
        std::cout << "客户端收到: " << msg << std::endl;
        message_received = true;
    });
    
    // 连接服务器
    assert(client.connect());
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    
    // 发送测试消息
    client.sendText("test message");
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    
    // 断开连接
    client.disconnect();
    server.stop();
    
    std::cout << "测试1完成" << std::endl << std::endl;
}

void test_multiple_clients() {
    std::cout << "测试2: 多客户端测试..." << std::endl;
    
    websocket::WebSocketServer server(8082);
    server.start();
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    
    int message_count = 0;
    server.setMessageHandler([&message_count](std::shared_ptr<websocket::WebSocketSession> session, const std::string& msg) {
        std::cout << "服务器收到消息 " << ++message_count << ": " << msg << std::endl;
        session->sendText("echo: " + msg);
    });
    
    // 创建多个客户端
    const int CLIENT_COUNT = 3;
    std::vector<std::unique_ptr<websocket::WebSocketClient>> clients;
    std::vector<int> received_counts(CLIENT_COUNT, 0);
    
    for (int i = 0; i < CLIENT_COUNT; ++i) {
        auto client = std::make_unique<websocket::WebSocketClient>("127.0.0.1", 8082);
        client->setMessageHandler([i, &received_counts](const std::string& msg) {
            std::cout << "客户端" << i << "收到: " << msg << std::endl;
            received_counts[i]++;
        });
        
        assert(client->connect());
        clients.push_back(std::move(client));
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }
    
    // 每个客户端发送消息
    for (int i = 0; i < CLIENT_COUNT; ++i) {
        clients[i]->sendText("hello from client " + std::to_string(i));
        std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }
    
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
    
    // 断开所有客户端
    for (auto& client : clients) {
        client->disconnect();
    }
    
    server.stop();
    std::cout << "测试2完成" << std::endl << std::endl;
}

int main() {
    std::cout << "开始WebSocket服务器测试..." << std::endl << std::endl;
    
    try {
        test_basic_communication();
        test_multiple_clients();
        
        std::cout << "所有测试通过!" << std::endl;
        return 0;
    } catch (const std::exception& e) {
        std::cerr << "测试失败: " << e.what() << std::endl;
        return 1;
    }
}